require 'optparse'
require 'sqlite3'
require 'rubygems'
require 'active_record'
require 'mechanize'
require 'uri'

$options = {}
parser = OptionParser.new("", 24) do |opts|
  opts.banner = "\nGoogle Scraper 1.0\nAuthor: Louis (Skype: louisprm)\n\n"

  opts.on("-o", "--output SQLITE3DB", "Output SQLite3 database file") do |v|
    $options[:output] = v
  end

  opts.on("-i", "--input FILE", "Input") do |v|
    $options[:input] = v
  end

  opts.on_tail('-h', '--help', 'Displays this help') do
		puts opts, "", help
    exit
	end
end

def help
  return <<-eos

GUIDELINE
-------------------------------------------------------
The scraper package includes two scripts

  1. scrape.rb: take a list of keywords and search against Google
                then write the results count to a local database
  2. export.rb: read the local database and generate the Excel/CSV output

For example:

  1. Search a list of keywords and write results to a main.db file
	   
        ruby scrape.rb --input=keywords.txt --output=main.db

  2. After the scraper script is done, run the export.rb script to read the main.db
     database and generate the Excel file search.xls

        ruby export.rb --input=main.db --output=/tmp/search.xls
 
eos
end

begin
  parser.parse!
rescue SystemExit => ex
  exit
rescue Exception => ex
  puts "\nERROR: #{ex.message}\n\nRun ruby scrape.rb -h for help\n\n"
  exit
end

if $options[:output].nil?
  puts "\nPlease specify output file: -o\n\n"
  exit
end

if $options[:input].nil?
  puts "\nPlease specify input file: -i\n\n"
  exit
end

$logger = Logger.new('/tmp/gscraper.log')

ActiveRecord::Base.establish_connection(
  adapter: 'sqlite3',
  database: $options[:output],
  timeout: 15000
)

require 'zlib'
class String
  def deflate
    Zlib.deflate(self)
  end

  def inflate
    Zlib.inflate(self)
  end

  def fix
    self.encode!('UTF-8', :undef => :replace, :invalid => :replace, :replace => "")
  end
end

class MySchema < ActiveRecord::Migration
  def change
    create_table :items do |t|
      t.text :keyword
      t.text :search_url
      t.integer :result_count
    end
  end
end

# initiate the database if not existed
MySchema.new.migrate(:change) unless File.exists?($options[:output])

class Item < ActiveRecord::Base; end

class Scrape
  BASE_URI = 'https://www.google.com/search?q='

  def initialize
    @a = Mechanize.new
    @a.agent.http.verify_mode = OpenSSL::SSL::VERIFY_NONE
  end

  def run
    IO.readlines($options[:input]).each do |keyword|
      search(keyword.strip) if keyword and keyword != ""
    end
  end

  def search(keyword)
    return if keyword.nil? or keyword.strip == ""
    search_url = "#{BASE_URI}#{URI.escape(keyword)}"

    if Item.exists?(search_url: search_url)
      $logger.info "Already exists: #{search_url}"
      return 
    end

    $logger.info "Searching: #{search_url}"

    ps = @a.get(search_url)

    p search_url
    
    i = Item.new
    i.keyword = keyword
    i.search_url = search_url
    if ps.parser.css('#resultStats').first
      i.result_count = ps.parser.css('#resultStats').first.text[/[0-9,]+/].gsub(',', '')
    else
      i.result_count = 0
    end

    i.save!

    $logger.info "Count: #{i.result_count}"
  end
end

# trap Ctrl-C
trap("SIGINT") { throw :ctrl_c }

catch :ctrl_c do
  e = Scrape.new
  e.run
end
